{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Welcome to Week 5 Day 1\n",
    "\n",
    "AutoGen AgentChat!\n",
    "\n",
    "This should look simple and familiar, because it has a lot in common with Crew and OpenAI Agents SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First concept: the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.models.ollama import OllamaChatCompletionClient\n",
    "ollamamodel_client = OllamaChatCompletionClient(model=\"llama3.2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second concept: The Message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.messages import TextMessage\n",
    "message = TextMessage(content=\"I'd like to go to London\", source=\"user\")\n",
    "message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Third concept: The Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "\n",
    "agent = AssistantAgent(\n",
    "    name=\"airline_agent\",\n",
    "    model_client=model_client,\n",
    "    system_message=\"You are a helpful assistant for an airline. You give short, humorous answers.\",\n",
    "    model_client_stream=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Put it all together with on_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core import CancellationToken\n",
    "\n",
    "response = await agent.on_messages([message], cancellation_token=CancellationToken())\n",
    "response.chat_message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's make a local database of ticket prices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sqlite3\n",
    "\n",
    "# Delete existing database file if it exists\n",
    "if os.path.exists(\"tickets.db\"):\n",
    "    os.remove(\"tickets.db\")\n",
    "\n",
    "# Create the database and the table\n",
    "conn = sqlite3.connect(\"tickets.db\")\n",
    "c = conn.cursor()\n",
    "c.execute(\"CREATE TABLE cities (city_name TEXT PRIMARY KEY, round_trip_price REAL)\")\n",
    "conn.commit()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Populate our database\n",
    "def save_city_price(city_name, round_trip_price):\n",
    "    conn = sqlite3.connect(\"tickets.db\")\n",
    "    c = conn.cursor()\n",
    "    c.execute(\"REPLACE INTO cities (city_name, round_trip_price) VALUES (?, ?)\", (city_name.lower(), round_trip_price))\n",
    "    conn.commit()\n",
    "    conn.close()\n",
    "\n",
    "# Some cities!\n",
    "save_city_price(\"London\", 299)\n",
    "save_city_price(\"Paris\", 399)\n",
    "save_city_price(\"Rome\", 499)\n",
    "save_city_price(\"Madrid\", 550)\n",
    "save_city_price(\"Barcelona\", 580)\n",
    "save_city_price(\"Berlin\", 525)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Method to get price for a city\n",
    "def get_city_price(city_name: str) -> float | None:\n",
    "    \"\"\" Get the roundtrip ticket price to travel to the city \"\"\"\n",
    "    conn = sqlite3.connect(\"tickets.db\")\n",
    "    c = conn.cursor()\n",
    "    c.execute(\"SELECT round_trip_price FROM cities WHERE city_name = ?\", (city_name.lower(),))\n",
    "    result = c.fetchone()\n",
    "    conn.close()\n",
    "    return result[0] if result else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_city_price(\"Rome\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_agentchat.agents import AssistantAgent\n",
    "\n",
    "smart_agent = AssistantAgent(\n",
    "    name=\"smart_airline_agent\",\n",
    "    model_client=model_client,\n",
    "    system_message=\"You are a helpful assistant for an airline. You give short, humorous answers, including the price of a roundtrip ticket.\",\n",
    "    model_client_stream=True,\n",
    "    tools=[get_city_price],\n",
    "    reflect_on_tool_use=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await smart_agent.on_messages([message], cancellation_token=CancellationToken())\n",
    "for inner_message in response.inner_messages:\n",
    "    print(inner_message.content)\n",
    "response.chat_message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
